from typing import Optional, List, Any, Mapping

from langchain.llms.base import LLM
from langchain.schema.callbacks.manager import CallbackManagerForLLMRun


class CustomLLM(LLM):
    """
    自定义LLM
    """
    n: int

    @property
    def _llm_type(self) -> str:
        return "custom"

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        if stop is not  None:
            raise ValueError("stop kwargs are not permitted.")
        return prompt[: self.n]

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """Get the indentifying parameters."""
        return {"n": self.n}


if __name__ == '__main__':
    llm = CustomLLM(n=10)
    llm("This is a foobar thing")
    print(llm)